/*
 * @Author       : Zhaoyu.Wu
 * @Date         : 2021-10-09 16:52
 * @LastEditTime : 2021-10-14 14:18
 * @LastEditors  : Zhaoyu.Wu
 * @Description  : AES library, include ECB, CBC, CFB, etc... 
 * 				   Used mbed TLS library
 * If you have any questions, email to zhaoyu.wu@diehl.com.
 */
/*******************************************************************************
 * Include files
 ******************************************************************************/
#include "aes.h"

#define AES_NULL ((void *)0)
#if defined(USING_AES_ASSERT)
#if defined(AES_ASSERT_PRINT)

/* 
 *define your private log output function if you need it
 */
// #define AES_PRINTF(fmt, ...) \
// 	LOG_E("AES", fmt, ##__VA_ARGS__)
#define AES_PRINTF(fmt, ...) \
	printf(fmt, ##__VA_ARGS__);

#define AES_ASSERT(condition)                         \
	do                                                \
	{                                                 \
		if (!(condition))                             \
		{                                             \
			AES_PRINTF("%s:%i: assert error: -%s\r\n",\
					   __FILE__,                      \
					   __LINE__,                      \
					   #condition);                   \
			return (-1);                              \
		}                                             \
	} while (0)

#else /* AES_ASSERT_PRINT */

#define AES_ASSERT(condition) \
	do                        \
	{                         \
		if (!(condition))     \
		{                     \
			return (-1);      \
		}                     \
	} while (0)

#endif /* AES_ASSERT_PRINT */

#else

#define AES_ASSERT(condition) \
	do                        \
	{                         \
		if (!(condition))	  \
		{                     \
		}                     \
	} while (0)

#endif /* USING_AES_ASSERT */

/*
 * AES-ECB block encryption/decryption
 */
#if defined (CRYPTO_TYPE_AES_ECB)
char aes_crypt_ecb(stc_aes_cfg_t *ctx)
{
	char ret = 0;
	if (ctx->enMode == AesMode_Encrypt)
	{
		while (ctx->u32ByteDataNum >= 16) //don't need padding
		{
			ret = _basic_aes_encrypt(ctx);
			AES_ASSERT(ret == 0);
			ctx->u32ByteDataNum -= 16;
			ctx->pu8InputByte += 16;
			ctx->pu8OutputByte += 16;
		}
		if (ctx->u32ByteDataNum > 0 && ctx->u32ByteDataNum < 16) //need padding
		{
			unsigned char pad_ary[16];
			unsigned char *input_addr, *output_addr;
			for (char i = 0; i < 16; i++) //padding the input data with parameter: padding
			{
				if (i < ctx->u32ByteDataNum)
				{
					pad_ary[i] = ctx->pu8InputByte[i];
				}
				else
				{
					pad_ary[i] = ctx->u8Padding;
				}
			}
			input_addr = ctx->pu8InputByte;
			output_addr = ctx->pu8OutputByte;
			ctx->pu8InputByte = pad_ary;
			ctx->pu8OutputByte = pad_ary;
			ret = _basic_aes_encrypt(ctx);
			AES_ASSERT(ret == 0);
			ctx->pu8InputByte = input_addr;
			ctx->pu8OutputByte = output_addr;
			for (char i = 0; i < 16; i++)
			{
				ctx->pu8OutputByte[i] = pad_ary[i];
			}
		}
		return 0;
	}
	else
	{
		while (ctx->u32ByteDataNum >= 16)
		{
			ret = _basic_aes_decrypt(ctx);
			AES_ASSERT(ret == 0);
			ctx->u32ByteDataNum -= 16;
			ctx->pu8InputByte += 16;
			ctx->pu8OutputByte += 16;
		}
		return 0;
	}
	return -1;
}
#endif	/* CRYPTO_TYPE_AES_ECB */

/*
 * AES-CBC block encryption/decryption
 */
#if defined(CRYPTO_TYPE_AES_CBC)
char aes_crypt_cbc(stc_aes_cfg_t *ctx)
{
	char ret = 0;
	if (ctx->enMode == AesMode_Encrypt)
	{
		while (ctx->u32ByteDataNum >= 16) //don't need padding
		{
			for (char i = 0; i < 16; i++)
			{
				ctx->pu8InputByte[i] = ctx->pu8InputByte[i] ^ ctx->pu8Ivec[i];
			}
			ret = _basic_aes_encrypt(ctx);
			AES_ASSERT(ret == 0);
			for (char i = 0; i < 16; i++)
			{
				ctx->pu8Ivec[i] = ctx->pu8OutputByte[i];
			}
			ctx->u32ByteDataNum -= 16;
			ctx->pu8InputByte += 16;
			ctx->pu8OutputByte += 16;
		}
		if (ctx->u32ByteDataNum > 0 && ctx->u32ByteDataNum < 16) //need padding
		{
			unsigned char pad_ary[16];
			unsigned char *input_addr, *output_addr;
			for (char i = 0; i < 16; i++) //padding the input data with parameter: padding, and xor with ivec
			{		
				if (i < ctx->u32ByteDataNum)
				{
					pad_ary[i] = ctx->pu8InputByte[i] ^ ctx->pu8Ivec[i];
				}
				else
				{
					pad_ary[i] = ctx->u8Padding ^ ctx->pu8Ivec[i];
				}	
			}
			input_addr = ctx->pu8InputByte;
			output_addr = ctx->pu8OutputByte;
			ctx->pu8InputByte = pad_ary;
			ctx->pu8OutputByte = pad_ary;
			ret = _basic_aes_encrypt(ctx);
			AES_ASSERT(ret == 0);
			ctx->pu8InputByte = input_addr;
			ctx->pu8OutputByte = output_addr;
			for (char i = 0; i < 16; i++)
			{
				ctx->pu8OutputByte[i] = pad_ary[i];
			}
		}
		return 0;
	}
	else
	{
		unsigned char temp_ivec[16];
		while (ctx->u32ByteDataNum >= 16)
		{
			for (char i = 0; i < 16; i++)
			{
				temp_ivec[i] = ctx->pu8InputByte[i];
			}
			ret = _basic_aes_decrypt(ctx);
			AES_ASSERT(ret == 0);
			for (char i = 0; i < 16; i++)
			{
				ctx->pu8OutputByte[i] = ctx->pu8OutputByte[i] ^ ctx->pu8Ivec[i];
				ctx->pu8Ivec[i] = temp_ivec[i];
			}
			ctx->u32ByteDataNum -= 16;
			ctx->pu8InputByte += 16;
			ctx->pu8OutputByte += 16;
		}
		return 0;
	}
	return -1;
}
#endif /* CRYPTO_TYPE_AES_CBC */


char aes_crypt(stc_aes_cfg_t* pstcAesCfg)
{
	AES_ASSERT(pstcAesCfg != AES_NULL);
    AES_ASSERT(pstcAesCfg->pu8InputByte != AES_NULL);
    AES_ASSERT(pstcAesCfg->pu8OutputByte != AES_NULL);
	AES_ASSERT(pstcAesCfg->pu8Key != AES_NULL);
	AES_ASSERT(pstcAesCfg->enMode == AesMode_Encrypt || 
				pstcAesCfg->enMode == AesMode_Decrypt);
	AES_ASSERT(pstcAesCfg->enKeyLen == AesKeyLen_128 ||
				pstcAesCfg->enKeyLen == AesKeyLen_192 ||
				pstcAesCfg->enKeyLen == AesKeyLen_256);
	if(pstcAesCfg->enMode == AesMode_Decrypt)
	{
		AES_ASSERT(pstcAesCfg->u32ByteDataNum % 16 == 0);
	}

#if defined(CRYPTO_TYPE_AES_ECB)
	if (pstcAesCfg->enCpoType == AesCpoType_Ecb)
	{
		return (aes_crypt_ecb(pstcAesCfg));
	}
#endif /* CRYPTO_TYPE_AES_ECB */

#if defined(CRYPTO_TYPE_AES_CBC)
	
	if (pstcAesCfg->enCpoType == AesCpoType_Cbc)
	{
		AES_ASSERT(pstcAesCfg->pu8Ivec != AES_NULL);
		return (aes_crypt_cbc(pstcAesCfg));
	}
#endif /* CRYPTO_TYPE_AES_CBC */

#if defined(CRYPTO_TYPE_AES_CFB)
	if (pstcAesCfg->enCpoType == AesCpoType_Cfb)
	{
		//TODO：AES_CFB
	}
#endif /* CRYPTO_TYPE_AES_CFB */

#if defined(CRYPTO_TYPE_AES_CTR)
	if (pstcAesCfg->enCpoType == AesCpoType_Ctr)
	{
		//TODO: AES_CTR
	}
#endif /* CRYPTO_TYPE_AES_CTR */

#if defined(CRYPTO_TYPE_AES_OFB)
	if (pstcAesCfg->enCpoType == AesCpoType_Ofb)
	{
		//TODO: AES_OFB
	}
#endif /* CRYPTO_TYPE_AES_OFB */
	return -2; //if enCpoType don't belong to en_aes_cpo_type_t, return -2
}




